#include <iostream>
#include <vector>
#include <thread>
#include <atomic>
#include "PcapLiveDeviceList.h"
#include "SystemUtils.h"

#include <pcap.h>
#include <chrono>//获取日期
#include <mutex>

#include <iostream>
#include <fstream>
#include <string>

#include <iomanip>//格式化当前日期
#include <sstream> 
#include <unordered_map>
#include <set>
#include <unordered_set>

#define START_SERVER 1  // 控制是否启用 HashTableManager

#if START_SERVER
#include <headers/HashTableManager.h>
#endif

enum IpValue {
    EXTRAHIGH,
    HIGH,
    MEDIUM,
    LOW,
    EXTRALOW
};
// 将枚举值转换为字符串
std::string ipValueToString(IpValue value) {
    switch (value) {
        case IpValue::EXTRALOW: return "很低";
        case IpValue::LOW: return "较低";
        case IpValue::MEDIUM: return "中等";
        case IpValue::HIGH: return "较高";
        case IpValue::EXTRAHIGH: return "很高";
        default: return "未知";
    }
}

/* 用来计算一个ip地址的价值度元组
 */
struct IpValueTuple
{
    double ipValueScore;
    IpValue ipValue; // ip地址价值度

    double GB; // 数据量
    int64_t packet_count; //包数

    std::set<uint32_t> session_ids; // 会话id集合
    int64_t session_count; //通联会话数

};


/**
 * A struct shared for collecting packet statistics
 */
struct PacketStats
{ 
    int ipv4PacketCount = 0;
    int ipv6PacketCount = 0;
    uint64_t tcpPacketCount = 0;
    int udpPacketCount = 0;
    int dnsPacketCount = 0;
    int httpPacketCount = 0;
    int sslPacketCount = 0;
    uint64_t unknownTcpPacketCount = 0;

    long long totalBits = 0; // 总比特数
    // long long totalPacketCounts = 0; // 总包数
    // long long totalSessionCount = 0; // 总会话数
    double totalTime = 0.0; // 总传输时间（秒）
    // double averageRate = 0.0;//平均速率
    std::chrono::time_point<std::chrono::system_clock> lastPacketTime; // 上一个数据包的时间
    bool firstPacket = true; // 标记是否为第一个数据包
    std::array<long long, 24 * 12> packetSize5minute = {0};    // 数据量（字节）
    std::array<long long, 24 * 12> averageRate5minute = {0};   // 平均速率（字节/秒）
    std::array<long long, 24 * 12> requestCount5minute = {0};  // 请求数量（基于流的统计）
    std::array<long long, 24 * 12> periodStartTimestamp = {0}; // 每个时间段的起始时间戳
    
    std::array<std::unordered_set<uint32_t>, 24 * 12> activeFlows;// 用于记录当前时间段的活跃流（五元组哈希值）
    std::vector<int64_t> packetCount5minute; // 每五分钟的包计数
    std::vector<int64_t> packetCounts; // 每两个小时的包计数
    std::vector<int64_t>packetCountsBits; // 每两个小时的包比特数
    std::unordered_map<std::string, IpValueTuple> Ip2IpValueTuple; // IP地址和价值度元组的映射
    std::mutex mutex; // 保护共享资源的互斥锁

    std::vector<int> tcpHeartbeatCounts; //每两个小时的异常包数

    #if START_SERVER
        HashTableManager manager;  // 条件编译: 启用服务器功能
    #endif

    // 统计指定端口的总流量（单位：字节）
    std::unordered_map<uint16_t, uint64_t> portTrafficMap;
    // 初始化
    PacketStats() : packetCount5minute(24*12,0),packetCounts(12, 0),packetCountsBits(12,0),lastPacketTime(std::chrono::system_clock::now()) {

    } 

     // 已知TCP端口白名单
    const std::unordered_set<uint16_t> KNOWN_TCP_PORTS = {
        80,    // HTTP
        443,   // HTTPS
        22,    // SSH
        25,    // SMTP
        53,    // DNS-over-TCP
        21,    // FTP
        3389,  // RDP
        3306   // MySQL
    };
    const size_t MAX_LINES = 10000; // 你可以根据需求调整

    void limit_csv_lines(const std::string& filename, size_t max_lines) {
        std::ifstream infile(filename);
        std::vector<std::string> lines;
        std::string line;
        while (std::getline(infile, line)) {
            lines.push_back(line);
        }
        infile.close();

        // 如果超出最大行数（包含表头），则删除多余的最旧数据
        if (lines.size() > max_lines + 1) { // +1 是表头
            // 保留表头和最后的max_lines条数据
            std::vector<std::string> new_lines;
            new_lines.push_back(lines[0]);
            new_lines.insert(new_lines.end(), lines.end() - max_lines, lines.end());
            std::ofstream outfile(filename, std::ios::trunc);
            for (auto& l : new_lines) outfile << l << "\n";
            outfile.close();
        }
    }

    // // 输出 CSV 文件
    // void outHeartbeatFile() {
    //     const std::string filename = "/home/njust/program/pcapplusplus-23.09-ubuntu-20.04-gcc-9.4.0-x86_64/backend-demo/stats/heartbeat.csv";
    //     bool needHeader = true;
    //     std::ifstream fin(filename);
    //     if (fin.is_open()) {
    //         std::string hdr;
    //         if (std::getline(fin, hdr))
    //         {
    //             if (hdr.rfind("DateTime", 0) == 0)
    //                 needHeader = false;
    //         }
    //         fin.close();
    //     }
    //     std::ofstream out(filename, std::ios::app);
    //     if (!out.is_open()) { std::cerr << "无法打开文件: " << filename << std::endl; return; }
    //     std::lock_guard<std::mutex> guard(mutex);
    //     if (needHeader) {
    //         out << "DateTime";
    //         for (size_t i = 0; i < tcpHeartbeatCounts.size(); ++i)
    //             out << ",Interval " << i*2 << " - " << i*2+2 << " hours";
    //         out << std::endl;
    //     }

    //     // 输出数据行
    //     auto now = std::chrono::system_clock::now();
    //     std::time_t t = std::chrono::system_clock::to_time_t(now);
    //     std::tm* tm = std::gmtime(&t);
    //     tm->tm_hour += 8; std::mktime(tm);
    //     char buf[64]; std::strftime(buf, sizeof(buf), "%Y-%m-%d %H:%M:%S", tm);
    //     out << buf;
    //     for (auto cnt : tcpHeartbeatCounts) out << "," << cnt;
    //     out << std::endl;
    //     out.close();
    //     limit_csv_lines(filename, MAX_HEARTBEAT_LINES);
    // }

    void consumePacketPortTraffic(pcpp::Packet& packet)
    {
        static const std::unordered_set<uint16_t> ports = {80, 443, 8080, 465, 21, 22};
        uint16_t dstPort = 0;
        size_t payloadLen = 0;

        // 只要能获取端口号和负载长度即可
        auto tcpLayer = packet.getLayerOfType<pcpp::TcpLayer>();
        if (tcpLayer) {
            dstPort = ntohs(tcpLayer->getTcpHeader()->portDst);
            payloadLen = tcpLayer->getLayerPayloadSize();
        }
        auto udpLayer = packet.getLayerOfType<pcpp::UdpLayer>();
        if (udpLayer) {
            dstPort = ntohs(udpLayer->getUdpHeader()->portDst);
            payloadLen = udpLayer->getLayerPayloadSize();
        }
        // 只要端口在目标集合里就统计
        if (ports.count(dstPort)) {
            std::lock_guard<std::mutex> lock(mutex);
            portTrafficMap[dstPort] += payloadLen;
        }
    }

    void printPortTraffic(std::ostream& out)
    {
        // 端口列表
        static const std::vector<uint16_t> ports = {80, 443, 8080, 465, 21, 22};

        // 输出表头
        out << "DateTime";
        for (auto port : ports)
            out << ",Port " << port << "(Bytes)";
        out << std::endl;

        // 获取东八区时间
        auto now = std::chrono::system_clock::now();
        std::time_t now_time_t = std::chrono::system_clock::to_time_t(now);
        std::tm* utc_time = std::gmtime(&now_time_t);
        utc_time->tm_hour += 8;
        std::mktime(utc_time);
        char buffer[100];
        std::strftime(buffer, sizeof(buffer), "%Y-%m-%d %H:%M:%S", utc_time);

        // 输出数据
        out << buffer;
        for (auto port : ports)
            out << "," << portTrafficMap[port];
        out << std::endl;
    }
    void out_port_file()
    {
        std::ofstream out("/home/njust/program/pcapplusplus-23.09-ubuntu-20.04-gcc-9.4.0-x86_64/backend-demo/stats/port_count.csv", std::ios::app); // 追加到文件末尾
        if (out.is_open())
        {
            // 加锁以确保线程安全
            std::lock_guard<std::mutex> guard(mutex);
            printPortTraffic(out); // 输出端口流量统计到文件

            out.close();
            // 限制最大行数
            limit_csv_lines("/home/njust/program/pcapplusplus-23.09-ubuntu-20.04-gcc-9.4.0-x86_64/backend-demo/stats/port_count.csv", MAX_LINES);
        }
        else
        {
            std::cerr << "无法打开文件: /home/njust/program/pcapplusplus-23.09-ubuntu-20.04-gcc-9.4.0-x86_64/backend-demo/stats/port_count.csv" << std::endl;
        }
    }

    void detectHeartbeat(pcpp::Packet& packet)
    {

    }
    void consumePacketProtocol(pcpp::Packet& packet)
    {        
        if (packet.isPacketOfType(pcpp::IPv4))
            ipv4PacketCount++;
        if (packet.isPacketOfType(pcpp::IPv6))
            ipv6PacketCount++;
        if (packet.isPacketOfType(pcpp::TCP)) {
            tcpPacketCount++;

            bool isKnown = false;
            
            // 检测应用层协议（HTTP/SSL等）
            if (packet.isPacketOfType(pcpp::HTTP) || 
                packet.isPacketOfType(pcpp::SSL)  || 
                packet.isPacketOfType(pcpp::FTP)) { // 可根据需要扩展
                isKnown = true;
            }

            // 若应用层未知，检查端口白名单
            if (!isKnown) {
                pcpp::TcpLayer* tcpLayer = packet.getLayerOfType<pcpp::TcpLayer>();
                if (tcpLayer != nullptr) {
                    uint16_t destPort = ntohs(tcpLayer->getTcpHeader()->portDst);
                    isKnown = KNOWN_TCP_PORTS.find(destPort) != KNOWN_TCP_PORTS.end();
                }
            }

            // 双重判断均为未知则计数
            if (!isKnown) {
                unknownTcpPacketCount++;
            }
        }
        if (packet.isPacketOfType(pcpp::UDP))
            udpPacketCount++;
        if (packet.isPacketOfType(pcpp::DNS))
            dnsPacketCount++;
        if (packet.isPacketOfType(pcpp::HTTP))
            httpPacketCount++;
        if (packet.isPacketOfType(pcpp::SSL))
            sslPacketCount++;
    }
    // 更新流量统计
    void updateSessionStats(pcpp::Packet& packet,const long long& totalPacketCounts)
    {
        #if START_SERVER
            uint32_t flowId = pcpp::hash5Tuple(&packet);
           long long SessionCount =  manager.updatePacketCount(flowId); //更新某个会话的总包数
        #endif
    }

    

    /**
     * Clear all stats
     */
    void clear()
    {
        std::lock_guard<std::mutex> guard(mutex);
        std::fill(packetCounts.begin(), packetCounts.end(), 0);
        ipv4PacketCount = ipv6PacketCount = tcpPacketCount = unknownTcpPacketCount = dnsPacketCount = httpPacketCount = sslPacketCount = 0;
    }

    /**
     * Collect stats from a packet
     */
    void consumePacket5minute(pcpp::Packet& packet)
    {
        auto now = std::chrono::system_clock::now();
        // 将当前时间转换为当前时间戳（秒）
        auto seconds = std::chrono::duration_cast<std::chrono::seconds>(now.time_since_epoch()).count();

        // 加上28800秒（8小时）以转换为CST时间
        seconds += 28800; // CST = UTC + 8 hours
        int index = (seconds / 300) % (24*12); // 每五分钟一个时间段
        // 获取当前时间段起始的时间戳
        int currentPeriodStart = (seconds / 300) * 300;
        std::lock_guard<std::mutex> guard(mutex);
        
        packetCount5minute[index]++;
        // 获取当前数据包的大小
        int packetSize = packet.getRawPacket()->getRawDataLen();
        // 累计数据包大小
        packetSize5minute[index] += packetSize;
        int elapsedSeconds = seconds - periodStartTimestamp[index];
        // int elapsedSeconds = seconds - currentPeriodStart;
        if (elapsedSeconds > 0)
        {
            averageRate5minute[index] = packetSize5minute[index] / elapsedSeconds;
        }
        else
        {
            averageRate5minute[index] = 0; // 无法计算速率时设置为 0
        }
        //计算哈希统计请求数
        // 如果时间段更新，则清理过期的流记录
        if (periodStartTimestamp[index] != currentPeriodStart)
        {
            periodStartTimestamp[index] = currentPeriodStart;
            activeFlows[index].clear(); // 清理上一个时间段的流记录
            requestCount5minute[index] = 0; // 重置请求数量
            averageRate5minute[index] = 0;  // 重置平均速率
        }

        // 使用 PcapPlusPlus 提供的 hash5Tuple 计算五元组的哈希值
        uint32_t flowId = pcpp::hash5Tuple(&packet);

        // 统计请求数量（基于流）
        if (activeFlows[index].find(flowId) == activeFlows[index].end())
        {
            // 如果当前流的哈希值不存在，则新增记录并增加请求计数
            activeFlows[index].insert(flowId);
            requestCount5minute[index]++;
        }

    }
    
    // 将时间戳转换为格式化字符串（例如：2024-12-24 14:30:00）
    std::string formatTimestamp(int timestamp)
    {
        std::time_t t = timestamp;
        std::tm* tm = std::gmtime(&t); // 转为 UTC 时间
        std::stringstream ss;
        ss << std::put_time(tm, "%Y-%m-%d %H:%M:%S");
        return ss.str();
    }

    /**
     * @brief 输出流量统计到CSV文件：[时间],[数据量KB],[包数量],[请求数],[平均速率KB/s]
     * 
     * 该函数每12秒执行一次，将5分钟粒度的流量统计数据写入CSV文件。
     * 处理逻辑：
     * 1. 读取现有CSV文件内容到内存(map结构)
     * 2. 更新内存中的统计数据
     * 3. 将更新后的数据写回文件
     * 
     * 特性：
     * - 线程安全：使用mutex保护共享资源
     * - 空文件处理：当没有数据时写入全0的一行
     * - 数据更新：保留历史数据，只更新变化的部分
     * 
     * @param filename 输出文件路径，默认为"./backend-demo/stats/trafficDetails.csv"
     * 
     * 输出格式：
     * Time Period,Data Size,Packet Count,Request Count,Average Rate
     * [时间],[数据量KB],[包数量],[请求数],[平均速率KB/s]
     * 
     * @note 文件路径是硬编码的，如需变更路径需要修改代码
     */
    void out_5min_file(const std::string& filename = "/home/njust/program/pcapplusplus-23.09-ubuntu-20.04-gcc-9.4.0-x86_64/backend-demo/stats/trafficDetails.csv")
    {
        std::lock_guard<std::mutex> guard(mutex);

        // 1. 读取现有文件内容到内存
        std::map<std::string, std::string> timePeriodData;
        std::ifstream in(filename);
        if (in.is_open()) {
            std::string line;
            if (std::getline(in, line)) { // 跳过表头
                while (std::getline(in, line)) {
                    size_t pos = line.find(',');
                    if (pos != std::string::npos) {
                        std::string timePeriod = line.substr(0, pos);
                        timePeriodData[timePeriod] = line;
                    }
                }
            }
            in.close();
        }

        // 2. 更新内存中的数据
        bool hasData = false;
        for (int i = 0; i < 24 * 12; i++) {
            if (packetCount5minute[i] > 0) {
                hasData = true;
                std::string timePeriod = formatTimestamp(periodStartTimestamp[i]);
                std::stringstream ss;
                ss << timePeriod << ","
                << packetSize5minute[i]/1024 << "KB,"
                << packetCount5minute[i] << ","
                << requestCount5minute[i] << ","
                << averageRate5minute[i]/1024 << "KB/s";
                timePeriodData[timePeriod] = ss.str();
            }
        }

        // 3. 写入文件
        std::ofstream out(filename, std::ios::trunc);
        if (!out.is_open()) {
            throw std::runtime_error("Unable to open file: " + filename);
        }
        
        out << "Time Period,Data Size,Packet Count,Request Count,Average Rate\n";
        
        if (hasData) {
            for (const auto& item : timePeriodData) {
                out << item.second << "\n";
            }
        } else {
            // 如果没有数据，写入全0的一行
            std::string currentTime = formatTimestamp(std::time(nullptr));
            out << currentTime << ",0KB,0,0,0KB/s\n";
        }
        out.close();
        // 限制最大行数
        limit_csv_lines("/home/njust/program/pcapplusplus-23.09-ubuntu-20.04-gcc-9.4.0-x86_64/backend-demo/stats/trafficDetails.csv", MAX_LINES);
    }

    /**
     * @brief 输出协议统计信息到CSV文件
     * 
     * 该函数将协议类型的统计信息(如TCP/UDP/ICMP等)追加写入到protocol_stats.csv文件
     * 每12秒执行一次，记录各协议的包数量占比情况
     * 
     * 特性：
     * - 线程安全：使用mutex保护共享资源
     * - 文件追加模式：不会覆盖已有数据
     * - 行数限制：自动维护文件最大行数(MAX_LINES)
     * 
     * 输出格式：
     * [时间戳],[协议类型1:数量],[协议类型2:数量],...
     * 
     * @note 文件路径硬编码为：/home/njust/program/pcapplusplus-23.09-ubuntu-20.04-gcc-9.4.0-x86_64/backend-demo/stats/protocol_stats.csv
     */
    void out_protocol_file()
    {
        auto filename = "/home/njust/program/pcapplusplus-23.09-ubuntu-20.04-gcc-9.4.0-x86_64/backend-demo/stats/protocol_stats.csv";
        std::ifstream in(filename);
        bool hasHead = false;
        if (in.is_open()) {
            std::string headerLine;
            if (std::getline(in, headerLine) && !headerLine.empty()) {
                hasHead = true;
            }
            in.close();
        }

        std::ofstream out(filename, std::ios::app); // 追加到文件末尾
        if (out.is_open())
        {
            // 加锁以确保线程安全
            std::lock_guard<std::mutex> guard(mutex);
            if (!hasHead) {
                out << "DateTime,IPv4 packet count,IPv6 packet count,TCP packet count,unknown TCP packet count,UDP packet count,DNS packet count,HTTP packet count,SSL packet count"<<std::endl;
            }
            auto now = std::chrono::system_clock::now();

            // 转换为 time_t
            std::time_t now_time_t = std::chrono::system_clock::to_time_t(now);

            // 转换为 UTC 时间
            std::tm* utc_time = std::gmtime(&now_time_t); // 获取 UTC 时间

            // 手动调整为东八区 (UTC+8)
            utc_time->tm_hour += 8;

            // 修正时间溢出（例如加 8 小时超过一天）
            std::mktime(utc_time); // 自动调整溢出部分，比如日期、月份
            char buffer[100];
            std::strftime(buffer, sizeof(buffer), "%Y-%m-%d %H:%M:%S", utc_time);
            out<<buffer<<","
                <<ipv4PacketCount<< ","
                <<ipv6PacketCount<< ","
                <<tcpPacketCount << ","
                <<unknownTcpPacketCount << ","
                <<udpPacketCount << ","
                <<dnsPacketCount << ","
                <<httpPacketCount<< ","
                <<sslPacketCount ;
            out << std::endl;
            
            out.close();
            // 限制最大行数
            limit_csv_lines("/home/njust/program/pcapplusplus-23.09-ubuntu-20.04-gcc-9.4.0-x86_64/backend-demo/stats/protocol_stats.csv", MAX_LINES);
        }
        else
        {
            std::cerr << "无法打开文件: /home/njust/program/pcapplusplus-23.09-ubuntu-20.04-gcc-9.4.0-x86_64/backend-demo/stats/protocol_stats.csv" <<std::endl;
        }
    }

    
    
    // 改成每两个小时去请求数据库统计 包数和比特量（行数和packet_len列和）写入到CSV文件
    // 平均速率同上，有没有求和的免重复计算
    // void consumePacket(pcpp::Packet& packet)
    // {
    //     auto now = std::chrono::system_clock::now();
    //     auto seconds = std::chrono::duration_cast<std::chrono::seconds>(now.time_since_epoch()).count();
    //     seconds += 28800; // CST = UTC + 8 hours
    //     int index = (seconds / 3600) % 24 / 2; // 每两个小时一个时间段
        
    //     {
    //         std::lock_guard<std::mutex> guard(mutex);
    //         packetCounts[index]++;
    //         totalPacketCounts++;
    //         // 获取数据包的字节大小
    //         size_t packetSizeBytes = packet.getRawPacket()->getRawDataLen();
    //         //std::cout << packetSizeBytes << std::endl;
    //         packetCountsBits[index] += packetSizeBytes*8;
    //         // 更新总比特数
    //         totalBits += packetSizeBytes*8; // 转换为bit
    //     }

    //      // 计算时间差并更新总时间
    //     if (firstPacket) 
    //     {
    //         firstPacket = false; // 第一个数据包处理,在处理第一个数据包时不更新 totalTime，确保只有在有数据包被处理后才计算速率。
    //         lastPacketTime = now; // 记录第一个数据包时间
    //     } 
    //     else 
    //     {
    //         double timeDiff = std::chrono::duration<double>(now - lastPacketTime).count(); // 计算时间差（秒）
    //         totalTime += timeDiff;
    //     }
    //     lastPacketTime = now; // 更新最后一个数据包时间
        
    //     // 更新平均速率
    //     if (totalTime > 0) 
    //     {
    //         averageRate = totalBits / totalTime; // 比特每秒
    //     }
    // }

   
    // // 获取当前日期字符串并且输出传输速率
    // std::string getCurrentDate()
    // {
    //     auto now = std::chrono::system_clock::now();
    //     auto time_t_now = std::chrono::system_clock::to_time_t(now);
    //     std::tm local_tm = *std::localtime(&time_t_now);
        
    //     // 格式化日期字符串
    //     std::ostringstream dateStream;
    //     dateStream << std::put_time(&local_tm, "%Y-%m-%d %H:%M:%S"); // 格式化为 "YYYY-MM-DD HH:MM:SS"

    //     // 输出平均传输速率
        
    //     // if (totalTime > 0)
    //     // {
    //     //     averageRate = totalBits / totalTime; // 比特每秒
    //     //     // std::cout << "totalBits: " << totalBits/1024/1024 << "mb" << std::endl;
    //     //     // std::cout << "totalTime: " << totalTime << "second" << std::endl;
    //     //     // std::cout << "Average transmission rate: " << averageRate/1024/1024 << " mb/second" << std::endl;
    //     // }
    //     // else 
    //     // {
    //     //     std::cout << "No data processed yet." << std::endl;
    //     // }
    //     return dateStream.str();
    // }

};